package com.demo.javaee;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * 多线程回滚, 只要有一个线程发生异常, 直接回滚
 * @author shenguangyang
 */
public class MultithreadingRollback {
    public static void main(String[] args) throws Exception {
        BossTask bossTask= new BossTask();

        WorkTask w1 = new MyTask01("A", bossTask, 3, true);
        WorkTask w2 = new MyTask01("B", bossTask, 2, false);
        WorkTask w3 = new MyTask02("C", bossTask, 1, true);

        bossTask.addWorkTask(w1);
        bossTask.addWorkTask(w2);
        bossTask.addWorkTask(w3);

        bossTask.waitDone();
        System.out.println("main end");
    }
}

enum TaskStatus {
    SUCCESS, FAIL
}

class BossTask {
    /** 等待工作线程超时 */
    private int waitTimeout = 5;
    private final List<WorkTask> workTasks;
    private volatile CountDownLatch workTaskLatch;

    public BossTask() {
        this.workTasks = new ArrayList<>();
    }

    public BossTask(int waitTimeout) {
        this();
        this.waitTimeout = waitTimeout;
    }

    public void addWorkTask(WorkTask workTask) {
        this.workTasks.add(workTask);
    }

    public boolean waitDone() {
        if (this.workTaskLatch == null) {
            synchronized (this) {
                if (this.workTaskLatch == null) {
                    this.workTaskLatch = new CountDownLatch(this.workTasks.size());
                    for (WorkTask task : workTasks) {
                        task.start();
                    }
                }
            }
        }
        // 等待超时
        try {
            return this.workTaskLatch.await(waitTimeout, TimeUnit.SECONDS);
        } catch (Exception e) {
            for (WorkTask task : workTasks) {
                task.cancel();
            }
            throw new RuntimeException(e);
        }
    }

    public void end(WorkTask workTask) {
        if (workTask.taskStatus == TaskStatus.FAIL) {
            cancel(workTask);
        }
        workTaskLatch.countDown();
    }

    public void cancel(WorkTask workTask) {
        for (WorkTask task : workTasks) {
            if (task != workTask) task.cancel();
        }
    }

}

abstract class WorkTask extends Thread {
    public TaskStatus taskStatus;
    private final BossTask bossTask;
    public volatile boolean canceling;
    public String threadName;

    public WorkTask(String threadName, BossTask bossTask) {
        this.taskStatus = TaskStatus.SUCCESS;
        this.canceling = false;
        this.bossTask = bossTask;
        this.threadName = threadName;
    }

    @Override
    public void run() {
        try {
            Thread.currentThread().setName(threadName);
            doRun();
        } catch (Exception e) {
            taskStatus = TaskStatus.FAIL;
            throw new RuntimeException(e);
        } finally {
            this.bossTask.end(this);
            this.bossTask.waitDone();
            if (canceling) {
                rollback();
            }
        }
    }

    /**
     * 完成任务
     */
    abstract void doRun();

    /**
     * 回滚
     */
    void rollback() {
        doRollback();
        this.canceling = false;
    }
    /**
     * 回滚
     */
    abstract void doRollback();

    public void cancel() {
        // 设置取消标志位
        this.canceling = true;
    }
}

/**
 * cpu密集型
 */
class MyTask01 extends WorkTask {
    boolean success;
    int timeInSeconds = 2;

    public MyTask01(String threadName, BossTask bossTask, int timeInSecond, boolean success) {
        super(threadName, bossTask);
        this.success = success;
        this.timeInSeconds = timeInSecond;
    }

    @Override
    void doRun() {
        System.out.println(Thread.currentThread().getName() + " start");
        try {
            TimeUnit.SECONDS.sleep(timeInSeconds);
            if (!this.success) {
                taskStatus = TaskStatus.FAIL;
            } else {
                System.out.println(Thread.currentThread().getName() + " success");
            }
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    void doRollback() {
        System.out.println(Thread.currentThread().getName() + " doRollback");
    }
}

/**
 * io密集型
 */
class MyTask02 extends WorkTask {
    int total;

    boolean success;
    int timeInSeconds = 2;

    public MyTask02(String threadName, BossTask bossTask, int timeInSecond, boolean success) {
        super(threadName, bossTask);
        this.success = success;
        this.timeInSeconds = timeInSecond;
    }

    @Override
    void doRun() {
        System.out.println(Thread.currentThread().getName() + " start");
        for (;;) {
            try {
                TimeUnit.SECONDS.sleep(1);
                if (total++ > 3) {
                    System.out.println(Thread.currentThread().getName() + " success");
                    break;
                }
                if (canceling) {
                    rollback();
                    break;
                }
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        if (canceling) {
            rollback();
        }
    }

    @Override
    void doRollback() {
        System.out.println(Thread.currentThread().getName() + " doRollback");
    }
}